package com.cyx.exercise.ollama;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIf;

import java.io.IOException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

import static org.junit.jupiter.api.Assertions.*;

/**
 * Ollama客户端测试类
 * 注意：这些测试需要Ollama服务运行并且安装了相应的模型
 */
class OllamaClientTest {

    private OllamaClient client;
    private static final String OLLAMA_BASE_URL = "http://localhost:11434";
    private static final String TEST_MODEL = "llama2";

    @BeforeEach
    void setUp() {
        client = new OllamaClient(OLLAMA_BASE_URL);
    }

    /**
     * 检查Ollama服务是否可用
     */
    private boolean isOllamaAvailable() {
        try {
            client.getModels();
            return true;
        } catch (Exception e) {
            System.out.println("Ollama服务不可用，跳过测试: " + e.getMessage());
            return false;
        }
    }

    @Test
    @EnabledIf("isOllamaAvailable")
    void testGetModels() {
        assertDoesNotThrow(() -> {
            String models = client.getModels();
            assertNotNull(models);
            assertFalse(models.trim().isEmpty());
            System.out.println("获取到的模型列表: " + models);
        });
    }

    @Test
    @EnabledIf("isOllamaAvailable")
    void testSyncGenerate() {
        assertDoesNotThrow(() -> {
            OllamaRequest request = new OllamaRequest(TEST_MODEL, "什么是1+1？请简短回答。");
            
            long startTime = System.currentTimeMillis();
            OllamaResponse response = client.generate(request);
            long endTime = System.currentTimeMillis();
            
            assertNotNull(response);
            assertNotNull(response.getResponse());
            assertFalse(response.getResponse().trim().isEmpty());
            assertTrue(response.isDone());
            
            System.out.println("同步生成测试:");
            System.out.println("问题: " + request.getPrompt());
            System.out.println("回答: " + response.getResponse());
            System.out.println("耗时: " + (endTime - startTime) + "ms");
        });
    }

    @Test
    @EnabledIf("isOllamaAvailable")
    void testAsyncGenerate() {
        assertDoesNotThrow(() -> {
            OllamaRequest request = new OllamaRequest(TEST_MODEL, "Java是什么？一句话回答。");
            
            long startTime = System.currentTimeMillis();
            CompletableFuture<OllamaResponse> future = client.generateAsync(request);
            
            // 等待异步结果，最多等待60秒
            OllamaResponse response = future.get(60, TimeUnit.SECONDS);
            long endTime = System.currentTimeMillis();
            
            assertNotNull(response);
            assertNotNull(response.getResponse());
            assertFalse(response.getResponse().trim().isEmpty());
            assertTrue(response.isDone());
            
            System.out.println("异步生成测试:");
            System.out.println("问题: " + request.getPrompt());
            System.out.println("回答: " + response.getResponse());
            System.out.println("异步耗时: " + (endTime - startTime) + "ms");
        });
    }

    @Test
    @EnabledIf("isOllamaAvailable")
    void testChat() {
        assertDoesNotThrow(() -> {
            OllamaChatRequest chatRequest = new OllamaChatRequest(TEST_MODEL);
            chatRequest.addSystemMessage("你是一个简洁的助手，用一句话回答问题。");
            chatRequest.addUserMessage("什么是Spring Boot？");
            
            OllamaChatResponse response = client.chat(chatRequest);
            
            assertNotNull(response);
            assertNotNull(response.getMessage());
            assertNotNull(response.getMessage().getContent());
            assertFalse(response.getMessage().getContent().trim().isEmpty());
            assertEquals("assistant", response.getMessage().getRole());
            assertTrue(response.isDone());
            
            System.out.println("聊天测试:");
            System.out.println("用户: 什么是Spring Boot？");
            System.out.println("助手: " + response.getMessage().getContent());
        });
    }

    @Test
    void testRequestValidation() {
        // 测试空模型名
        assertThrows(Exception.class, () -> {
            OllamaRequest request = new OllamaRequest("", "测试问题");
            client.generate(request);
        });

        // 测试null模型名
        assertThrows(Exception.class, () -> {
            OllamaRequest request = new OllamaRequest(null, "测试问题");
            client.generate(request);
        });
    }

    @Test
    void testInvalidModel() {
        // 测试不存在的模型
        if (isOllamaAvailable()) {
            assertThrows(IOException.class, () -> {
                OllamaRequest request = new OllamaRequest("non-existent-model-12345", "测试问题");
                client.generate(request);
            });
        }
    }

    @Test
    void testChatRequestBuilder() {
        OllamaChatRequest request = new OllamaChatRequest(TEST_MODEL);
        
        // 测试添加消息
        request.addSystemMessage("系统消息");
        request.addUserMessage("用户消息");
        request.addAssistantMessage("助手消息");
        
        assertEquals(3, request.getMessages().size());
        assertEquals("system", request.getMessages().get(0).getRole());
        assertEquals("系统消息", request.getMessages().get(0).getContent());
        assertEquals("user", request.getMessages().get(1).getRole());
        assertEquals("用户消息", request.getMessages().get(1).getContent());
        assertEquals("assistant", request.getMessages().get(2).getRole());
        assertEquals("助手消息", request.getMessages().get(2).getContent());
    }

    @Test
    void testResponseObjects() {
        // 测试OllamaResponse
        OllamaResponse response = new OllamaResponse("测试响应", true);
        assertEquals("测试响应", response.getResponse());
        assertTrue(response.isDone());
        
        // 测试OllamaChatMessage
        OllamaChatMessage message = new OllamaChatMessage("user", "测试内容");
        assertEquals("user", message.getRole());
        assertEquals("测试内容", message.getContent());
        
        // 测试OllamaChatResponse
        OllamaChatResponse chatResponse = new OllamaChatResponse(message, true);
        assertEquals(message, chatResponse.getMessage());
        assertTrue(chatResponse.isDone());
    }

    @Test
    void testClientConfiguration() {
        // 测试不同的基础URL
        OllamaClient customClient = new OllamaClient("http://example.com:11434");
        assertNotNull(customClient);
        
        // 测试URL末尾斜杠处理
        OllamaClient clientWithSlash = new OllamaClient("http://localhost:11434/");
        assertNotNull(clientWithSlash);
    }

    @Test
    @EnabledIf("isOllamaAvailable")
    void testStreamGenerate() throws InterruptedException {
        if (!isOllamaAvailable()) {
            return;
        }

        OllamaRequest request = new OllamaRequest(TEST_MODEL, "数到3，每个数字一行。");
        StringBuilder result = new StringBuilder();
        
        // 使用CountDownLatch等待流式响应完成
        java.util.concurrent.CountDownLatch latch = new java.util.concurrent.CountDownLatch(1);
        
        client.generateStream(request,
            chunk -> {
                result.append(chunk);
                System.out.print(chunk);
            },
            error -> {
                System.err.println("流式生成错误: " + error.getMessage());
                latch.countDown();
            }
        );
        
        // 等待一段时间让流式响应完成
        Thread.sleep(5000);
        
        System.out.println("\n流式生成完整结果: " + result.toString());
        assertFalse(result.toString().trim().isEmpty());
    }
}
